import pandas as pd
import statsmodels.formula.api as sm

from cat_rfs.code.utils.table import Table

SETTINGS = {'main': {'beta_var': 'beta_sheet',
                     'yield_var': 'sheet_yield_usd',
                     'table_file': 'table6'},
            'trace': {'beta_var': 'beta_act',
                      'yield_var': 'act_yield_usd',
                      'table_file': 'table6_trace'},
            'loss': {'beta_var': 'beta_loss',
                     'yield_var': 'sheet_yield_usd',
                     'table_file': 'table6_loss'}}


def print_table6(df, output='console', setting='main'):
    tname = f"{SETTINGS[setting]['table_file']}_corr"
    beta = SETTINGS[setting]['beta_var']

    df = df.copy()
    df = df.dropna(subset=['turnover'])  # Turnover is na for bonds that never had a TRACE observation
    df = df[(df['date'] >= '12/31/2004') & (df['date'] <= '12/31/2018')]

    df['CUSIP6'] = df['CUSIP9'].str[:6]

    df = df.sort_values(['date', 'CUSIP9']).reset_index(drop=True).copy()
    df['time_to_maturity'] = (df['current_maturity'] - df['date']).dt.days / 365
    assert (df['time_to_maturity'] > 0).all()

    df['age'] = (df['date'] - df['pricing_date']).dt.days / 365
    assert (df['age'] >= 0).all()

    varnames = ['$\\hat{\\beta}$', 'Turnover', 'N trades', 'Trd intvl', 'Size', 'Age', 'Maturity', 'Dealer spread']
    corr = df[[beta, 'turnover', 'n_trades', 'trading_interval',
               'size', 'age', 'time_to_maturity', 'dealer_spread']].corr()
    corr.insert(0, '', varnames)

    T = Table(9, width='auto', justs=['l', 'r', 'r', 'r', 'r', 'r', 'r', 'r', 'r'],
              caption='Correlation of beta and liquidity measures',
              label=tname)

    T.add_panel(corr, formatting="%.2f", headers=[''] + varnames)

    if output == 'paper':
        T.print_table(f'cat_rfs/output/tables/{tname}.tex')
    else:
        T.print_table()

    yield_ = SETTINGS[setting]['yield_var']
    tname = f"{SETTINGS[setting]['table_file']}_pricing"

    df = df.copy()
    df['er'] = df[yield_] - df['expected_loss'] / 100 - df['rf'] / 100
    df = df.dropna(subset=[beta, 'er'])

    df['er_m'] = df['er'] * df['size']
    df2 = df.groupby('date', as_index=False)[['er_m', 'size']].sum()
    df2['er_m'] = df2['er_m'] / df2['size']
    df = df.drop(columns=['er_m'])
    df = df.merge(df2[['date', 'er_m']], on='date')

    df['er_model'] = df[beta] * df['er_m']

    df['cluster'] = df['CUSIP6'] + df['issue_num'].astype(str)
    dfc = df.copy()
    data = pd.DataFrame([['$\\lambda_0$', '', '$\\lambda_{cat}$', '',
                         '$\\lambda_{liq}$', '']],
                        columns=['alpha', 'alpha_t', 'lambda', 'lambda_se',
                                 'lambda_liq', 'lambda_liq_se'])

    other_stats = [['$N$', '$R^2$']]
    for var in ('turnover', 'n_trades', 'trading_interval', 'size', 'age', 'time_to_maturity', 'dealer_spread'):

        df = dfc.copy()

        if var == 'dealer_spread':
            df = df[df['date'] >= '12/1/2008'].copy()

        df = df.dropna(subset=[var])
        df[var] = (df[var] - df[var].mean()) / df[var].std()

        res = []
        for t in df['date'].unique():
            df2 = df[(df['date'] == t)].reset_index()
            ols = sm.ols(formula=f'er ~ {beta} + {var}', data=df2).fit()

            a = ols.params['Intercept'] * 100

            la = ols.params[beta] * 100

            lal = ols.params[2] * 100

            res.append([pd.to_datetime(t).year, a, la, lal, ols.rsquared, ols.nobs])

        res_df = pd.DataFrame(res, columns=['year', 'alpha', 'lambda_', 'lambda_liq', 'Rsquared', 'N'])

        lags = 4
        ols = sm.ols('alpha ~ 1', data=res_df).fit(cov_type='HAC', cov_kwds={'maxlags': lags}, use_t=True)
        a = ols.params['Intercept']
        a_se = ols.bse['Intercept']

        ols = sm.ols('lambda_ ~ 1', data=res_df).fit(cov_type='HAC', cov_kwds={'maxlags': lags}, use_t=True)
        la = ols.params['Intercept']
        la_se = ols.bse['Intercept']

        ols = sm.ols('lambda_liq ~ 1', data=res_df).fit(cov_type='HAC', cov_kwds={'maxlags': lags}, use_t=True)
        lal = ols.params['Intercept']
        lal_se = ols.bse['Intercept']

        N = res_df.shape[0]
        R2 = res_df['Rsquared'].mean()

        data = data.append(pd.DataFrame([[a, a_se, la, la_se, lal, lal_se]],
                                        columns=['alpha', 'alpha_t', 'lambda', 'lambda_se',
                                                 'lambda_liq', 'lambda_liq_se']), ignore_index=True, sort=False)

        other_stats.append([N, R2])

    T = Table(8, width='auto', justs=['l', 'd', 'd', 'd', 'd', 'd', 'd', 'd'],
              caption='Pricing of catastrophe market risk and liquidity',
              label=tname, footer=None)
    T.add_panel(data.T, supheader_cols=[1, (2, 8)], supheaders=['', 'Liquidity measure ($LIQ$)'],
                headers=['', 'Turnover', 'N trades', 'Trd intvl', 'Size', 'Age', 'Maturity', 'Dealer spread'],
                regression_specs={'stars': (0.1, 0.05, 0.01), 'dof': (N-1)})
    T.add_panel(other_stats)

    if output == 'paper':
        T.print_table(f'cat_rfs/output/tables/{tname}.tex')
    else:
        T.print_table()

    pass
